package com.sui.bigdata.rtcadmin.util;

import com.alibaba.fastjson.JSON;
import com.feidee.fdcommon.configuration.CustomConfiguration;
import com.google.common.collect.Maps;
import com.sui.bigdata.rtcadmin.constant.JobConstant;
import com.sui.bigdata.rtcadmin.dto.JobConfigDto;
import org.apache.commons.lang.StringUtils;
import org.apache.flink.configuration.*;
import org.apache.flink.yarn.configuration.YarnConfigOptions;

import java.util.Map;

/**
 * @author YongChen
 * @date 2019/11/15 16:28
 * @description
 * @email yong_chen@sui.com
 */
public class ConfUtils {
        private static final String APPLICATION_TYPE = "Sui Flink";
        private static final String DEFAULT_TASKMANAGER_MEMORY = "768m";
        private static final String DEFAULT_JOBMANAGER_MEMORY = "768m";
        private static final String DEFAULT_TASKSLOTS = "1";


    public static Map<String, String> confPropUtil(JobConfigDto jobConfigDto, Configuration flinkConfig) {
        String taskManagerMemory = jobConfigDto.getTaskManagerMemory() <= 0 ? DEFAULT_TASKMANAGER_MEMORY : jobConfigDto.getTaskManagerMemory() + "m";
        String jobManagerMemory = jobConfigDto.getJobManagerMemory() <= 0 ? DEFAULT_JOBMANAGER_MEMORY : jobConfigDto.getJobManagerMemory() + "m";
        String taskSlots = jobConfigDto.getSlots() <= 0 ? DEFAULT_TASKSLOTS : jobConfigDto.getSlots() + "";

        Map<String, String> confProp = Maps.newHashMap();
        confProp.put(YarnConfigOptions.APPLICATION_TYPE.key(), APPLICATION_TYPE);
        confProp.put(YarnConfigOptions.APPLICATION_NAME.key(), jobConfigDto.getJobName());

        confProp.put(TaskManagerOptions.NUM_TASK_SLOTS.key(), taskSlots);
        confProp.put(TaskManagerOptions.TOTAL_PROCESS_MEMORY.key(), taskManagerMemory);
        confProp.put(JobManagerOptions.TOTAL_PROCESS_MEMORY.key(), jobManagerMemory);


        if(jobConfigDto.getParallelism() > 0){
            confProp.put(CoreOptions.DEFAULT_PARALLELISM.key(), jobConfigDto.getParallelism() + "");
        }
        if (StringUtils.isNotBlank(jobConfigDto.getQueue())) {
            confProp.put(YarnConfigOptions.APPLICATION_QUEUE.key(), jobConfigDto.getQueue());
        }
        if (StringUtils.isNotBlank(jobConfigDto.getExtendParam()) && !"nan".equals(jobConfigDto.getExtendParam())) {
            confProp.putAll(JSON.parseObject(jobConfigDto.getExtendParam(), Map.class));
        }
        if(StringUtils.isBlank(confProp.get(CheckpointingOptions.CHECKPOINTS_DIRECTORY.key()))){
            confProp.put(CheckpointingOptions.CHECKPOINTS_DIRECTORY.key(),
                    flinkConfig.getString(CheckpointingOptions.CHECKPOINTS_DIRECTORY) + "/" + jobConfigDto.getJobName());
        }

        return confProp;
    }
}
